package net.jcip.examples.blocking;

import java.math.BigInteger;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

public class Memoizer<A, V> implements Computable<A, V> {

    private final static Long TIMEOUT_EXPIRATION = 1000l;

    private final Computable<A, V> computer;
    private final ConcurrentMap<A, Future<V>> cache = new ConcurrentHashMap<A, Future<V>>();
    private final ConcurrentMap<A, Long> expirationTime = new ConcurrentHashMap<A, Long>();

    public Memoizer(Computable<A, V> computer) {
        this.computer = computer;
    }

    @Override
    public V compute(final A arg) throws InterruptedException  {
        while (true) {
            Long time = expirationTime.get(arg);
            if (time != null && (time - System.currentTimeMillis() > TIMEOUT_EXPIRATION)) {
                expirationTime.remove(arg);
                cache.remove(arg);
            }
            Future<V> future = cache.get(arg);
            if (future == null) {
                Callable<V> eval = new Callable<V>() {
                    @Override
                    public V call() throws Exception {
                        return computer.compute(arg);
                    }
                };
                FutureTask<V> futureTask = new FutureTask<V>(eval);
                future = cache.putIfAbsent(arg, future);
                if (future == null) {
                    future = futureTask;
                    futureTask.run();
                    expirationTime.putIfAbsent(arg, System.currentTimeMillis());
                }
            }
            try {
                return future.get();
            } catch (CancellationException e) {
                cache.remove(arg, future);
            } catch (ExecutionException e) {
                cache.remove(arg, future);
                throw new RuntimeException(e);
            }
        }
    }
}

class ExpensiveComputation implements Computable<String, BigInteger> {

    @Override
    public BigInteger compute(String arg) throws InterruptedException {
        return new BigInteger(arg);
    }

}

class TestComputation implements Computable<String, String> {

    private final Map<String, Integer> countCall = new HashMap<String, Integer>();

    @Override
    public String compute(String arg) throws InterruptedException, ExecutionException {
        synchronized (countCall) {
            if (countCall.containsKey(arg)) {
                countCall.put(arg, countCall.get(arg) + 1);
            } else {
                countCall.put(arg, 0);
            }
        }
        return arg;
   }
}
